package com.wj.tree;

import java.util.NoSuchElementException;

public class RedBlackBST<Key extends Comparable<Key>, Value> {

    private static final boolean RED   = true;
    private static final boolean BLACK = false;

    //红黑树的根节点
    private Node root;

    private class Node {
        private Key key;
        private Value val;
        private Node left, right;  // 左子树和右子树
        private boolean color;     // 链接颜色
        private int size;          // 子树的大小

        public Node(Key key, Value val, boolean color, int size) {
            this.key = key;
            this.val = val;
            this.color = color;
            this.size = size;
        }
    }

    private boolean isRed(Node x) {
        if (x == null) return false;
        return x.color == RED;
    }

    private int size(Node x) {
        if (x == null) return 0;
        return x.size;
    }

    public int size() {
        return size(root);
    }

    public boolean isEmpty() {
        return root == null;
    }

    /**
     * @author wen.jie
     * @date 2021/8/25 14:16
     * 根据指定key查询value
     */
    public Value get(Key key) {
        if (key == null) throw new IllegalArgumentException("argument to get() is null");
        return get(root, key);
    }

    /**
     * @author wen.jie
     * @date 2021/8/25 14:16
     * 该树中是否包含指定的key
     */
    public boolean contains(Key key) {
        return get(key) != null;
    }

    private Value get(Node x, Key key) {
        while (x != null) {
            int cmp = key.compareTo(x.key);
            if      (cmp < 0) x = x.left;
            else if (cmp > 0) x = x.right;
            else              return x.val;
        }
        return null;
    }

    public void put(Key key, Value val) {
        if (key == null) throw new IllegalArgumentException("first argument to put() is null");
        root = put(root, key, val);
        //树的根节点总是黑色
        root.color = BLACK;
    }

    private Node put(Node h, Key key, Value val) {
        if (h == null) return new Node(key, val, RED, 1);

        //这里与二叉查找树代码一样
        int cmp = key.compareTo(h.key);
        if      (cmp < 0) h.left  = put(h.left,  key, val);
        else if (cmp > 0) h.right = put(h.right, key, val);
        else              h.val   = val;
        //红黑树平衡化
        return balance(h);
    }

    // 树的平衡化
    private Node balance(Node h) {
        //左旋:右链接为红,左连接为黑
        if (isRed(h.right) && !isRed(h.left))    h = rotateLeft(h);
        //右旋:两条连续的左链接
        if (isRed(h.left) && isRed(h.left.left)) h = rotateRight(h);
        //节点的左链接和右链接都是红链接,发生颜色变换
        if (isRed(h.left) && isRed(h.right))     flipColors(h);
        h.size = size(h.left) + size(h.right) + 1;
        return h;
    }

    //右旋
    private Node rotateRight(Node h) {
        assert (h != null) && isRed(h.left);
        Node x = h.left;
        h.left = x.right;
        x.right = h;
        x.color = x.right.color;
        x.right.color = RED;
        x.size = h.size;
        h.size = size(h.left) + size(h.right) + 1;
        return x;
    }

    //左旋
    private Node rotateLeft(Node h) {
        assert (h != null) && isRed(h.right);
        Node x = h.right;
        h.right = x.left;
        x.left = h;
        x.color = x.left.color;
        x.left.color = RED;
        x.size = h.size;
        h.size = size(h.left) + size(h.right) + 1;
        return x;
    }

    // 颜色变换
    private void flipColors(Node h) {
        h.color = !h.color;
        h.left.color = !h.left.color;
        h.right.color = !h.right.color;
    }

    public void  delMin(){
        if (isEmpty()){
            throw new NoSuchElementException("树为空");
        }
        if(!isRed(root.left) && !isRed(root.right)){
            root.color = RED;   // 如果根节点的左右子节点是2-节点，我们要先根设为红的，这样才能进行后面的moveRedLeft操作，因为左孩子要从根节点借一个
        }
        root = delMin(root);
        root.color = BLACK;  // 借完以后，我们将根节点的颜色复原
    }
    private Node delMin(Node node) {
        if (node.left==null){
            return null;
        }
        if (!isRed(node.left) && !isRed(node.left.left)){
            // node的左节点如果是2-节点，则按上面的方法编程3-节点或是临时4-节点
            node=moveRedLeft(node);
        }
        node.left=delMin(node.left);
        return balance(node); //   平衡临时组成的4-节点
    }

    private Node moveRedLeft(Node node) {
        /**
         * 因为我们规定红链接只能在左，
         * 因此当前节点的左右子节点都是2-节点，这时候我们就需要通过颜色转换，将这三个节点合并在一起
         */
        flipColors(node);
        //如果兄弟节点为2-节点的话，那么到上一步就结束了

        if(isRed(node.right.left)){     // 而如果兄弟节点不是2-节点的话，我们就需要通过旋转从兄弟节点借一个过来
            node.right = rotateRight(node.right);
            node = rotateLeft(node);
            //  因为条件2要求我们只向兄弟节点借一个，
            //  而一开始从父节点那里借了一个，因此需要还一个给父节点
            flipColors(node);
        }
        return node;

    }

    public void delMax(){
        if (isEmpty()){
            throw new NoSuchElementException("树为空");
        }
        if (!isRed(root.left) && !isRed(root.right)){
            root.color=RED;
        }
        root=delMax(root);
        root.color=BLACK;
    }

    private Node delMax(Node node) {
        if(isRed(node.left)){

            node = rotateRight(node);
        }
        if (node.right==null){
            return null;
        }
        if (!isRed(node.right) && !isRed(node.right.left)){
            node=moveRedRight(node);
        }
        node.right=delMax(node.right);
        return balance(node);
    }

    private Node moveRedRight(Node h) {
        flipColors(h);
        if (isRed(h.left.left)) {
            h = rotateRight(h);
            flipColors(h);
        }
        return h;
    }

    public void delete(Key key) {
        if (key == null) throw new IllegalArgumentException("argument to delete() is null");
        if (!contains(key)) return;

        if (!isRed(root.left) && !isRed(root.right))
            root.color = RED;

        root = delete(root, key);
        if (!isEmpty()) root.color = BLACK;
    }
    private Node delete(Node node, Key key) {

        if (key.compareTo(node.key) < 0)  {
            //key在左子树，按删除最小值一样删除
            if (!isRed(node.left) && !isRed(node.left.left))
                node = moveRedLeft(node);
            node.left = delete(node.left, key);
        }
        else {//key在右子树，按删除最大值一样删除
            if (isRed(node.left))
                node = rotateRight(node);
            //需首先判断待删除的节点是否在树底，否则下一步的node.right.left会出现空指针
            if (key.compareTo(node.key) == 0 && (node.right == null))
                return null;
            //递归在右子树中删除
            if (!isRed(node.right) && !isRed(node.right.left))
                node = moveRedRight(node);
            //如果待删除的节点不在树底
            if (key.compareTo(node.key) == 0) {
                //像二叉查找树一样删除，从node右子树找到最小的节点放在当前位置，然后再将该最小节点从右子树删除
                Node x = min(node.right);
                node.key = x.key;
                node.val = x.val;

                node.right = delMin(node.right);
            }
            else {
                //key仍然大于该node，继续在右子树递归
                node.right = delete(node.right, key);
            }
        }
        return balance(node);
    }

    public Key min() {
        if (isEmpty()) throw new NoSuchElementException("calls min() with empty symbol table");
        return min(root).key;
    }

    private Node min(Node x) {
        if (x.left == null) return x;
        else                return min(x.left);
    }


}
